# -*- coding: utf-8 -*- 
import urllib, urllib2 
import re, sys, os
import datetime, time
import traceback
import urlparse
import binascii, hashlib
import cStringIO as StringIO
import threading
import argparse
import paramiko
import workerpool

from lxml       import etree
from urlparse   import urljoin
from termcolor  import cprint
from pymongo    import Connection
stringify = etree.XPath("string()")
LOCAL_PATH          = '/home/hoangnamhai/HarvestedData/gallery'
MONGO_SERVER        = '27.0.12.106'   
#MONGO_SERVER        = 'beta.mana.vn'   
MONGO_PORT          = 27017
DATABASE            = 'gallery'
PREFIX              = '/uploads/gallery' 
os.umask(0000)
start = 0
def getArgs():
    parser  = argparse.ArgumentParser()
    parser.add_argument('-e', '--run', '-run', '-copy', action='store', dest='copy_file', help='Copy file to local => server.')
    args    = parser.parse_args()
    return  args.copy_file

def createSSHClient(server, port, user, passwd):
    try:
        client  = paramiko.SSHClient()
        client.load_system_host_keys()
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        client.connect(server, port, user, passwd)

        return client
    except:
        traceback.print_exc()

def rexists(sftp, path):
    try:
        sftp.stat(path)
        return True
    except:
        return False

flgCopy = getArgs()
ssh = None; sftp = None
if flgCopy!=None:
    ssh     = createSSHClient('mana.vn', 22, 'daudm', '')
    sftp    = ssh.open_sftp()
    if ssh==None:
        print 'Could not connect via paramiko .'
        forceQuit()

class GetImg():
    def __init__(self, server, port, database):
        CONNECT   = Connection(server, port)
        self.DB   = CONNECT[database]
        self.data = []

    def getPathWithName(self, src):
        ''' Hàm lấy tên ảnh qua xử lý md5(src)
            Giá trị trả về là chuỗi mã hóa md5 và đường dẫn (ex: a/b/c/d)
        '''
        try:
            if type(src).__name__ == 'unicode': src = src.encode('utf-8')
            md5str = hashlib.md5(src).hexdigest()
            path = '{0}/{1}/{2}/{3}'.format(md5str[0], md5str[1], md5str[2], md5str[3])
            return md5str, path
        except:
            traceback.print_exc()
            return None, None

    def saveImage(self, src, prefix, localPath):
        ''' Hàm lưu ảnh với src, prefix, local_path
            Giá trị trả về là : 
                +) result (0 chưa tồn tại, 1: ảnh đã tồn tại)
                +) source (prefix/path/file_name) Ex: /uploads/xkcn/a/b/c/d/abcd.....jpg
                +) img_name: tên ảnh là md5(src).jpg
                +) size : Dung lượng của ảnh   
        '''
        try:
            if src==None or src=='': return None, None, None, None
            if type(src).__name__ == 'unicode': src = src.encode('utf-8')
            img_name, path = self.getPathWithName(src)
            if img_name==None: return None, None, None, None
            img_ext      = os.path.splitext(src)[1]
            img_ext      = img_ext.lower()
            if len(img_ext)>4:  
                img_ext = img_ext[:4]
                if img_ext=='.jpe': img_ext = '.jpeg'
            if len(img_ext)<=1: img_ext = '.jpg' 
            preg         = re.compile('jpg|png|gif|png')
            m            = preg.search(img_ext)
            if not m: img_ext = '.jpg'
            path_file    = '{0}/{1}/{2}{3}'.format(localPath, path, img_name, img_ext)
            path_dir     = '{0}/{1}'.format(localPath, path)
            source       = '{0}/{1}/{2}{3}'.format(prefix, path, img_name, img_ext)
            file_name    = '{0}{1}'.format(img_name, img_ext)
            flag = False; retry = 0; maxRetry = 3; result = 0; size = 0
            while not flag and retry < maxRetry:
                try:
                    if not os.path.isfile(path_file):
                        if not os.path.exists(path_dir): os.makedirs(path_dir, 0777)
                        location, response = urllib.urlretrieve(src, path_file)
                        if not response.has_key('Content-Type'):
                            os.unlink(location); result = None
                            raise Exception('%s: content-type not exists!' % src)
                        else:
                            if not re.search('image', response['Content-Type']): 
                                os.unlink(location); result = None
                                raise Exception('%s: file is not image' % src)
                        size = os.path.getsize(location)
                        if long(response['Content-Length'])!=size:
                            os.unlink(location); result = None
                            raise Exception('%s: size of file not equal file response' % src)
                        if size > 0:
                            cprint('Success !', 'green'); break
                        else: os.unlink(location); result = None; break
                    else:
                        size = os.path.getsize(path_file)
                        cprint('File is exists !', 'red'); result = 1; break
                except:
                    traceback.print_exc()
                    retry += 1
            return result, source, file_name, size
        except:
            traceback.print_exc()
            return None, None, None, None

    def saveImageWithSCP(self, src, prefix, localPath, ssh=None, sftp=None):
        try:
            if src==None or src=='' or ssh==None or sftp==None: return None, None, None, None
            if type(src).__name__ == 'unicode': src = src.encode('utf-8')
            img_name, path = self.getPathWithName(src)
            if img_name==None: return None, None, None, None
            img_ext      = os.path.splitext(src)[1]
            img_ext      = img_ext.lower()
            if len(img_ext)>4:  
                img_ext = img_ext[:4]
                if img_ext=='.jpe': img_ext = '.jpeg'
            if len(img_ext)<=1: img_ext = '.jpg' 
            preg         = re.compile('jpg|png|gif|png')
            m            = preg.search(img_ext)
            if not m: img_ext = '.jpg'
            path_file    = '{0}/{1}/{2}{3}'.format(localPath, path, img_name, img_ext)
            path_dir     = '{0}/{1}'.format(localPath, path)
            source       = '{0}/{1}/{2}{3}'.format(prefix, path, img_name, img_ext)
            file_name    = '{0}.jpg'.format(img_name)
            flag = False; retry = 0; maxRetry = 3; result = 0; size = 0
            while not flag and retry < maxRetry:
                try:
                    if not os.path.isfile(path_file):
                        if not os.path.exists(path_dir): os.makedirs(path_dir, 0777)
                        location, response = urllib.urlretrieve(src, path_file)
                        if not response.has_key('Content-Type'):
                            os.unlink(location); result = None
                            raise Exception('%s: content-type not exists!' % src)
                        else:
                            if not re.search('image', response['Content-Type']): 
                                os.unlink(location); result = None
                                raise Exception('%s: file is not image' % src)
                        size = os.path.getsize(location)
                        if long(response['Content-Length'])!=size:
                            os.unlink(location); result = None
                            raise Exception('%s: size of file not equal file response' % src)
                        if size > 0:
                            stdin, stdout, stderr   = ssh.exec_command("mkdir -p -m 0777 {0}".format(path_dir))
                            remote_file = path_file
                            info = rexists(sftp, remote_file)
                            if not info:
                                flgRun = False
                                for i in range(0, 3):
                                    try:
                                        sftp.put(path_file, remote_file)
                                        print 'Copy file: ', path_file, ' => ', remote_file
                                        flgRun = True; break
                                    except:
                                       continue 
                                if flgRun:
                                    cprint('Upload thành công!', 'green'); break
                                else: 
                                    cprint('Không upload được ảnh lên', 'red'); result = None; break
                            else: cprint('Ảnh đã tồn tại trên server!', 'red')
                        else: os.unlink(location); result = None
                    else:
                        size = os.path.getsize(path_file)
                        stdin, stdout, stderr   = ssh.exec_command("mkdir -p -m 0777 {0}".format(path_dir))
                        remote_file = path_file
                        info = rexists(sftp, remote_file)
                        if not info:
                            flgRun = False
                            for i in range(0, 3):
                                try:
                                    sftp.put(path_file, remote_file)
                                    print 'Copy file: ', path_file, ' => ', remote_file
                                    flgRun = True; break
                                except:
                                   continue 
                            if flgRun:
                                cprint('Upload ảnh thành công!', 'green')
                            else: 
                                cprint('Không upload được ảnh lên', 'red'); result = None; break
                        else: cprint('Ảnh đã tồn tại trên server!', 'red')
                        cprint('Ảnh đã tồn tại !', 'red'); result = 1; break
                except:
                    traceback.print_exc()
                    retry += 1
            return result, source, file_name, size
        except:
            traceback.print_exc()
            return None, None, None, None

    def processObj(self, obj):
        try:
            url = obj['url']
            caption = obj['caption']
            result = None; source = file_name = ''; size = 0
            if flgCopy!=None:
                result, source, file_name, size = self.saveImageWithSCP(url, PREFIX, LOCAL_PATH, ssh, sftp)
                if result==None:
                    for i in range(0, 3):
                        result, source, file_name, size = self.saveImageWithSCP(url, PREFIX, LOCAL_PATH, ssh, sftp)
                        if result!=None: break
            else:
                result, source, file_name, size = self.saveImage(url, PREFIX, LOCAL_PATH)
                if result==None:
                    for i in range(0, 3):
                        result, source, file_name, size = self.saveImage(url, PREFIX, LOCAL_PATH)
                        if result!=None: break
            if result!=None:
               doc = {'url': url, 'caption': caption, 'source': source, 'name': file_name, 'size': size}
               print doc
               return doc
        except:
            traceback.print_exc()

    def processItem(self, data):
        try:
            if data==None: return
            listObj = []
            listImg = []
            items = data['data']
            for item in items:
                obj = {'url': item['url'], 'caption': item['caption']}
#                listObj.append(obj)
                result = self.processObj(obj) 
                if result!=None and result!=False:
                    listImg.append(result)
#            pool = workerpool.WorkerPool(size=1)
#            pool.map(self.processObj, listObj)
#            pool.shutdown(); pool.wait()
            if len(listImg) > 0: 
#                print listImg
                obj_id = data['_id']
                collection = self.DB['album']
                collection.update({'_id': obj_id}, {'$set': {'data': listImg, 'is_active': True}})
        except:
            traceback.print_exc()

    def processAll(self):
        try:
            collection = self.DB['album']
            rows = collection.find({'is_active': False}, {'data': 1})
            for row in rows:
                self.processItem(row); self.data = []
        except:
            traceback.print_exc()

def forceQuit():
    print 'Finished. ', datetime.datetime.now()
    pid = os.getpid(); os._exit(1); os.kill(pid, 9)

def timeOut():        
    while True:
        delta = time.time() - start
        if delta > 1800:
            print 'Dừng chương trình vì vượt quá thời gian chạy.', datetime.datetime.now()
            print 'process timeout {0}'.format(delta)
            pid = os.getpid(); os._exit(1); os.kill(pid, 9)
        time.sleep(30)

if __name__ == '__main__':
    try:
        cprint('start tool get images !!', 'yellow')
        print 'Started. ', datetime.datetime.now()
        start = time.time() 
        timeout = threading.Thread(target=timeOut).start()
        crawler = GetImg(MONGO_SERVER, MONGO_PORT, DATABASE)
        crawler.processAll()
        if flgCopy!=None: ssh.close()
        forceQuit()
    except:
        traceback.print_exc()
